# models/model_strategy_manager.py
from typing import Dict, Any, Optional, Callable, List
import torch
import asyncio
from enum import Enum

from .lightweight_ppo import MultiHeadPPO
try:
    from utils.online_model_service import OnlineModelService
except ImportError:
    # 备用导入路径
    from online_model_service import OnlineModelService

class ModelType(Enum):
    """模型类型枚举"""
    LOCAL_LIGHTWEIGHT = "local_lightweight"  # 本地轻量级模型
    ONLINE_POWERFUL = "online_powerful"      # 在线高性能模型
    HYBRID = "hybrid"                        # 混合模型

class ModelStrategyManager:
    """模型策略管理器 - 管理不同组件使用不同模型的策略"""
    
    def __init__(self, 
                 local_model: Optional[MultiHeadPPO] = None,
                 online_model_service: Optional[OnlineModelService] = None):
        self.local_model = local_model
        self.online_model_service = online_model_service
        self.model_strategies = {}
        self.component_contexts = {}  # 存储各组件的上下文
        
    def register_strategy(self, component_name: str, model_type: ModelType):
        """注册组件的模型策略"""
        self.model_strategies[component_name] = model_type
        
    def get_model_type(self, component_name: str) -> Optional[ModelType]:
        """获取组件的模型类型"""
        return self.model_strategies.get(component_name)
    
    async def execute_decision(
        self, 
        component_name: str, 
        state: Any, 
        decision_context: Dict[str, Any] = None
    ) -> Dict[str, Any]:
        """根据策略执行决策"""
        model_type = self.get_model_type(component_name)
        
        if not model_type:
            raise ValueError(f"未找到组件 {component_name} 的模型策略")
        
        if model_type == ModelType.LOCAL_LIGHTWEIGHT:
            return await self._execute_local_decision(component_name, state, decision_context)
        elif model_type == ModelType.ONLINE_POWERFUL:
            return await self._execute_online_decision(component_name, state, decision_context)
        elif model_type == ModelType.HYBRID:
            return await self._execute_hybrid_decision(component_name, state, decision_context)
        else:
            raise ValueError(f"未知的模型类型: {model_type}")
    
    async def _execute_local_decision(
        self, 
        component_name: str, 
        state: Any, 
        decision_context: Dict[str, Any] = None
    ) -> Dict[str, Any]:
        """执行本地轻量级模型决策"""
        if not self.local_model:
            raise ValueError("本地模型未初始化")
        
        try:
            # 使用本地PPO模型进行决策
            actions, log_probs, values, entropies = self.local_model.compute_actions(
                {component_name: [state]},  # 简化输入格式
                [component_name] # 因子类型
            )
            
            return {
                "action": actions.get(component_name, {}),
                "log_probs": log_probs.get(component_name, {}),
                "value": values.get(component_name, {}).get('value', 0),
                "model_used": "local_lightweight",
                "success": True
            }
        except Exception as e:
            print(f"本地模型决策失败: {e}")
            return {
                "action": {},
                "error": str(e),
                "model_used": "local_lightweight",
                "success": False
            }
    
    async def _execute_online_decision(
        self, 
        component_name: str, 
        state: Any, 
        decision_context: Dict[str, Any] = None
    ) -> Dict[str, Any]:
        """执行在线高性能模型决策"""
        if not self.online_model_service:
            raise ValueError("在线模型服务未初始化")
        
        try:
            # 构建决策提示
            prompt = self._build_decision_prompt(component_name, state, decision_context)
            
            # 使用在线模型进行决策
            result = await self.online_model_service.call_model(
                prompt=prompt,
                system_message=self._get_system_message(component_name),
                max_tokens=500,
                temperature=0.3  # 决策任务使用较低温度以获得更一致的结果
            )
            
            if result and result["success"]:
                # 解析在线模型的输出
                parsed_result = self._parse_online_decision(result["content"], state)
                
                return {
                    "action": parsed_result,
                    "raw_response": result["content"],
                    "model_used": "online_powerful",
                    "model_name": result["model"],
                    "success": True
                }
            else:
                # 在线模型失败，返回错误信息
                return {
                    "action": {},
                    "error": "在线模型调用失败",
                    "model_used": "online_powerful",
                    "success": False
                }
                
        except Exception as e:
            print(f"在线模型决策失败: {e}")
            return {
                "action": {},
                "error": str(e),
                "model_used": "online_powerful",
                "success": False
            }
    
    async def _execute_hybrid_decision(
        self, 
        component_name: str, 
        state: Any, 
        decision_context: Dict[str, Any] = None
    ) -> Dict[str, Any]:
        """执行混合模型决策 - 结合本地和在线模型"""
        if not self.local_model and not self.online_model_service:
            raise ValueError("本地模型和在线模型服务都未初始化")
        
        # 根据决策复杂度选择模型
        decision_complexity = self._assess_decision_complexity(state, decision_context)
        
        if decision_complexity > 0.7:  # 复杂决策使用在线模型
            online_result = await self._execute_online_decision(component_name, state, decision_context)
            if online_result["success"]:
                return online_result
            else:
                # 在线模型失败，降级到本地模型
                return await self._execute_local_decision(component_name, state, decision_context)
        else:  # 简单决策使用本地模型
            local_result = await self._execute_local_decision(component_name, state, decision_context)
            if local_result["success"]:
                return local_result
            else:
                # 本地模型失败，升级到在线模型
                return await self._execute_online_decision(component_name, state, decision_context)
    
    def _build_decision_prompt(
        self, 
        component_name: str, 
        state: Any, 
        decision_context: Dict[str, Any] = None
    ) -> str:
        """构建决策提示"""
        prompt_parts = []
        
        # 添加组件角色描述
        if "admin" in component_name.lower() or "trainer" in component_name.lower():
            role_description = "你是训练系统的管理员，负责制定高级策略和复杂决策。"
        else:
            role_description = "你是一个智能体，需要根据当前状态做出合适的决策。"
        
        prompt_parts.append(role_description)
        
        # 添加状态信息
        prompt_parts.append(f"当前状态: {state}")
        
        # 添加上下文信息
        if decision_context:
            prompt_parts.append(f"决策上下文: {decision_context}")
        
        # 添加决策要求
        prompt_parts.append(
            "请根据以上信息，提供一个具体的决策或行动建议。"
            "如果涉及数值，请提供具体的数字。"
            "请确保决策是可执行的。"
        )
        
        return "\n".join(prompt_parts)
    
    def _get_system_message(self, component_name: str) -> str:
        """获取系统消息"""
        if "admin" in component_name.lower() or "trainer" in component_name.lower():
            return (
                "你是一个高级AI训练管理员，负责制定复杂策略和优化训练流程。"
                "请提供专业、准确且可执行的决策建议。"
                "对于技术决策，请考虑性能、效率和稳定性。"
            )
        else:
            return (
                "你是一个AI智能体，需要根据当前环境状态做出合适的决策。"
                "请提供具体、可执行的行动建议。"
            )
    
    def _parse_online_decision(self, response: str, original_state: Any) -> Dict[str, Any]:
        """解析在线模型的决策输出"""
        # 简单的解析逻辑，实际应用中可能需要更复杂的解析
        try:
            # 尝试解析为JSON格式
            import json
            import re
            
            # 提取JSON部分（如果有）
            json_match = re.search(r'\{.*\}', response, re.DOTALL)
            if json_match:
                json_str = json_match.group()
                return json.loads(json_str)
            
            # 如果不是JSON，返回基本结构
            return {
                "decision": response[:200],  # 截取前200字符作为决策
                "confidence": 0.8,  # 假设置信度
                "parsed_from": "text"
            }
        except:
            # 解析失败时返回基本结构
            return {
                "decision": str(response),
                "confidence": 0.5,
                "parsed_from": "fallback"
            }
    
    def _assess_decision_complexity(
        self, 
        state: Any, 
        decision_context: Dict[str, Any] = None
    ) -> float:
        """评估决策复杂度 (0.0-1.0)"""
        complexity = 0.0
        
        # 根据状态复杂度评估
        if hasattr(state, '__len__') and len(state) > 50:  # 状态向量较长
            complexity += 0.3
        elif hasattr(state, '__len__') and len(state) > 20:
            complexity += 0.2
        
        # 根据上下文复杂度评估
        if decision_context and len(decision_context) > 5:
            complexity += 0.3
            
        # 根据决策选项数量评估
        if decision_context and 'available_actions' in decision_context:
            action_count = len(decision_context['available_actions'])
            if action_count > 10:
                complexity += 0.4
            elif action_count > 5:
                complexity += 0.2
        
        return min(complexity, 1.0)  # 限制在0-1范围内
    
    def update_context(self, component_name: str, context: Dict[str, Any]):
        """更新组件上下文"""
        if component_name not in self.component_contexts:
            self.component_contexts[component_name] = []
        
        self.component_contexts[component_name].append(context)
        
        # 限制上下文历史长度
        if len(self.component_contexts[component_name]) > 50:
            self.component_contexts[component_name] = self.component_contexts[component_name][-50:]
    
    def get_context_history(self, component_name: str, limit: int = 10) -> List[Dict[str, Any]]:
        """获取组件上下文历史"""
        if component_name in self.component_contexts:
            return self.component_contexts[component_name][-limit:]
        return []

# 便利函数：创建默认策略管理器
def create_default_strategy_manager(
    factor_types: list,
    online_api_key: Optional[str] = None
) -> ModelStrategyManager:
    """创建默认的策略管理器"""
    # 初始化本地模型
    local_model = MultiHeadPPO(factor_types=factor_types)
    
    # 初始化在线模型服务
    online_service = None
    if online_api_key:
        online_service = OnlineModelService(api_key=online_api_key)
    
    # 创建策略管理器
    manager = ModelStrategyManager(
        local_model=local_model,
        online_model_service=online_service
    )
    
    # 注册默认策略
    # 管理员组件使用在线模型
    manager.register_strategy("training_controller", ModelType.ONLINE_POWERFUL)
    manager.register_strategy("admin_system", ModelType.ONLINE_POWERFUL)
    manager.register_strategy("system_optimizer", ModelType.ONLINE_POWERFUL)
    
    # 智能体使用本地模型
    for factor_type in factor_types:
        manager.register_strategy(factor_type.name.lower(), ModelType.LOCAL_LIGHTWEIGHT)
    
    return manager
